from pydantic import BaseModel, Field
from typing import List, Optional, Union
yolo_labels="81:gloves;82:small tissue pack;83:large tissue box;84:used tissue."
face_labels="0:lulixin;1:geyiheng;2:aijiahao;3:chengyiling.(since the voice recog may not transcribe the name correctly, you should allow some variance when I say find someone, and choose the correct person id.)"
################################################################################
# LEG TOOLS
################################################################################
class ForwardModel(BaseModel):
    distance: int=Field(...,description="distance to move forward (cm)")
class TurnModel(BaseModel):
    direction: int=Field(...,description="direction to turn.-135(left) to 135(right)")
################################################################################
# ARM TOOLS
################################################################################
class ArmMoveModel(BaseModel):
    x: float=Field(...,description="x coordinate of the target")
    y: float=Field(...,description="y coordinate of the target")
    z: float=Field(...,description="z coordinate of the target")
    releaseOnPos: bool=Field(default=False,description="if true, the robot will release the object on moving to the target position")
    catchOnPos: bool=Field(default=False,description="if true, the robot will catch the object on moving to the target position")
################################################################################
# CAMERA TOOLS
################################################################################
class CaptureModel(BaseModel):
    yaw_degree: float=Field(default=0,description="yaw degree of the camera.ranging from -135(left) to 135(right)")
    pitch_degree: float=Field(default=0,description="pitch degree of the camera.ranging from -30 to 90.If you want to see the ground, for example, you should set it to -10")
################################################################################
# VOICE TOOLS
################################################################################
class SpeakModel(BaseModel):
    content: str=Field(...,description="content to speak")
################################################################################
# SCHEDULE TOOLS
################################################################################
class SetScheduleModel(BaseModel):
    jobtype: str=Field(...,description="cron or interval.cron means a one-time task,interval means a task that repeats at a certain interval.")
    task: str=Field(...,description="description for the task for TaskAgent to execute in the future")
    hour: int=Field(...,description="hour to set for interval or cron")
    minute: int=Field(...,description="minute to set for interval or cron")
    second: int=Field(...,description="second to set for interval or cron")
################################################################################
# MAIN TOOL
################################################################################
class FindWalkModel(BaseModel):
    searchWith:str=Field(...,description="searchWith can be 'face' or 'yolo' or 'vllm'.If you want to find a person, you should set it to 'face'.If you want to find an object, you should set it to 'yolo'")
    label: Union[int,str]=Field(...,description=f"if searchWith vllm, write STRING label here.If searchWith face or yolo, write INT label here.{face_labels} . {yolo_labels}")
    walk: bool=Field(default=True,description="if true, the robot will walk to the target")
    catch: bool=Field(default=False,description="if true, the robot will catch the target")
    putToPack: bool=Field(default=False,description="if true, the robot will put the catched object to the back pack")
class GrabThingsModel(BaseModel):
    walk: bool=Field(default=False,description="if true, the robot will walk to the target")
class SwitchWheelLegModel(BaseModel):
    wheelOrLeg: int=Field(...,description="0:wheel,1:leg")
class WheelMoveModel(BaseModel):
    dist:int = Field(...,description="the distance(meter) of forwarding")
AGENT_TOOLS=[
    # ForwardModel,
    #          TurnModel,
    #          FindWalkModel,
    #          CaptureModel,
    #          SpeakModel,
    #          SetScheduleModel,
    # GrabThingsModel,
    SwitchWheelLegModel
    ]

